#!/usr/bin/env bash

DIR=$(realpath $0) && DIR=${DIR%/*}
cd $DIR
set -ex

[ "$UID" -eq 0 ] || exec "$0" "$@"

if ! command -v nvcc &> /dev/null
then

echo "卸载旧的 NVIDIA 驱动和 CUDA..."
apt-get remove --purge -y '^nvidia-.*' '^cuda-.*' || true
apt-get autoremove -y
apt-get autoclean

echo "安装推荐驱动..."
apt update
recommended_driver=$(ubuntu-drivers devices | grep recommended | awk '{print $3}')
apt install -y "$recommended_driver"

echo "未找到nvcc，正在尝试安装CUDA工具包..."
echo "deb http://security.ubuntu.com/ubuntu focal-security main" | tee /etc/apt/sources.list.d/focal-security.list
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
add-apt-repository -y "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"

apt-get update
apt-get install -y libssl1.1 cuda cuda-toolkit

echo CUDA工具包安装完成
echo 需要重启
fi

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export CUDA_ROOT=/usr/local/cuda
echo "正在尝试重新编译项目..."
cargo install --path mistralrs-server --features "cuda flash-attn cudnn"
